# for semi-final
# tang-xg@qq.com 20200403
'''
target_cross(目标路口):
    需要提交预测的路口。（训练集和测试集中都没有任何交通流的信息）
relate_cross(相关路口):
    跟目标路口挨着或在同一条路上的，且在训练集和测试集中有交通流数据的路口(目标路口与相关路口的对应字典)。
pred_cross(预测路口):
    跟目标路口挨着或在同一条路上的，且在训练集和测试集中有交通流数据的路口(所有目标路口的相关路口的集合)。
similar_cross(相似路口):【后续没用】
    跟目标路口挨着或在同一条路上的路口。
'''

import numpy as np
import pandas as pd 
import os
import pickle as pkl

os.chdir("/home/ubuntu/txg/traf_qingdao/semi-final")

pwd = os.getcwd()
train_path = pwd + "/train/"
test_path = pwd + "/test_user/"

# =============================================================================
# 1, run get_GPS.py,   strore GPS in crossName['lat','lng']
# =============================================================================
# os.system("python getGPS_v3_baidu.py")


# =============================================================================
# 2, load data of data.pkl with GPS
# =============================================================================
[flow,relate_cross,crossName,target_cross,pred_cross,similar_cross,roadNet,submitList0,data_tr]=pkl.load(open(pwd+"/data.pkl",'rb'))


# =============================================================================
# 3, submitList add timeStamp, then save and reload
# =============================================================================
submitList = pd.read_csv(pwd+"/submit_example.csv",index_col=None)
from datetime import datetime
submitList['timeStamp'] = np.zeros([len(submitList),1])
for k in range(len(submitList)):
    timeStr = '2019-' + str(submitList.loc[k,'date'])+' '+str(submitList.loc[k,'timeBegin'])+':00'
    time_ = datetime.strptime(timeStr,'%Y-%m-%d %H:%M:%S')
    #timeStr = timeType.strftime('%Y-%m-%d %H:%M:%S')
    submitList.loc[k,'timeStamp'] = time_
# save and reload
pkl.dump([flow,relate_cross,crossName,target_cross,pred_cross,similar_cross,roadNet,submitList,data_tr],open(pwd+"/data.pkl",'wb'))
[flow,relate_cross,crossName,target_cross,pred_cross,similar_cross,roadNet,submitList,data_tr]=pkl.load(open(pwd+"/data.pkl",'rb'))


# =============================================================================
# 4, add flow2, add direction1, direction2 
# =============================================================================
train_list = []
for fileName in os.listdir(train_path):
    if fileName[:5] == 'train':
        train_list.append(fileName)
train_list.sort(key=lambda x:int(x.split('-')[-1].split('.')[0]))

test_list = []
for fileName in os.listdir(test_path):
    if fileName[:5] == 'test_':
        test_list.append(fileName)
test_list.sort(key=lambda x:int(x.split('-')[-1].split('.')[0]))

rng = pd.date_range(start='2019-09-01 07:00:00',end='2019-09-25 18:55:00',freq = '5min')


# flow2
flow2 = pd.DataFrame(np.zeros([len(rng),len(pred_cross)]),index=rng)
flow2.columns = np.sort(list(pred_cross))

for fileName in train_list:
    print(fileName)
    tmp = pd.read_csv(train_path+fileName,index_col=None)
    for k in range(len(tmp)):
        crossroadID = tmp.loc[k,'crossroadID']
        if crossroadID in pred_cross:
            timestamp = tmp.loc[k,'timestamp']
            timestamp = timestamp[:-2] + '00'
            minute = int( np.floor(int(timestamp[-5:-3]) / 5) * 5 )
            timestamp = timestamp[:-5] + str(minute) + timestamp[-3:]
            flow2.loc[timestamp, crossroadID] = flow2.loc[timestamp, crossroadID] + 1
        
for fileName in test_list:
    print(fileName)
    tmp = pd.read_csv(test_path+fileName,index_col=None)
    for k in range(len(tmp)):
        crossroadID = tmp.loc[k,'crossroadID']
        if crossroadID in pred_cross:
            timestamp = tmp.loc[k,'timestamp']
            timestamp = timestamp[:-2] + '00'
            minute = int( np.floor(int(timestamp[-5:-3]) / 5) * 5 )
            timestamp = timestamp[:-5] + str(minute) + timestamp[-3:]
            flow2.loc[timestamp, crossroadID] = flow2.loc[timestamp, crossroadID] + 1

pkl.dump([flow2],open(pwd+"/flow2.pkl",'wb'))

# direction1
direction1 = pd.DataFrame(np.zeros([len(rng),len(pred_cross)]),index=rng)
direction1.columns = np.sort(list(pred_cross))
for fileName in train_list:
    print(fileName)
    tmp = pd.read_csv(train_path+fileName,index_col=None)
    for k in range(len(tmp)):
        crossroadID = tmp.loc[k,'crossroadID']
        if crossroadID in pred_cross:
            timestamp = tmp.loc[k,'timestamp']
            timestamp = timestamp[:-2] + '00'
            minute = int( np.floor(int(timestamp[-5:-3]) / 5) * 5 )
            timestamp = timestamp[:-5] + str(minute) + timestamp[-3:]
            direction1.loc[timestamp, crossroadID] = direction1.loc[timestamp, crossroadID] + 1
for fileName in test_list:
    print(fileName)
    tmp = pd.read_csv(test_path+fileName,index_col=None)
    for k in range(len(tmp)):
        crossroadID = tmp.loc[k,'crossroadID']
        if crossroadID in pred_cross:
            timestamp = tmp.loc[k,'timestamp']
            timestamp = timestamp[:-2] + '00'
            minute = int( np.floor(int(timestamp[-5:-3]) / 5) * 5 )
            timestamp = timestamp[:-5] + str(minute) + timestamp[-3:]
            direction1.loc[timestamp, crossroadID] = direction1.loc[timestamp, crossroadID] + 1
pkl.dump([direction1],open(pwd+"/direction1.pkl",'wb'))


# direction2
direction2 = pd.DataFrame(np.zeros([len(rng),len(pred_cross)]),index=rng)
direction2.columns = np.sort(list(pred_cross))
for fileName in train_list:
    print(fileName)
    tmp = pd.read_csv(train_path+fileName,index_col=None)
    for k in range(len(tmp)):
        crossroadID = tmp.loc[k,'crossroadID']
        if crossroadID in pred_cross:
            timestamp = tmp.loc[k,'timestamp']
            timestamp = timestamp[:-2] + '00'
            minute = int( np.floor(int(timestamp[-5:-3]) / 5) * 5 )
            timestamp = timestamp[:-5] + str(minute) + timestamp[-3:]
            direction2.loc[timestamp, crossroadID] = direction2.loc[timestamp, crossroadID] + 1
for fileName in test_list:
    print(fileName)
    tmp = pd.read_csv(test_path+fileName,index_col=None)
    for k in range(len(tmp)):
        crossroadID = tmp.loc[k,'crossroadID']
        if crossroadID in pred_cross:
            timestamp = tmp.loc[k,'timestamp']
            timestamp = timestamp[:-2] + '00'
            minute = int( np.floor(int(timestamp[-5:-3]) / 5) * 5 )
            timestamp = timestamp[:-5] + str(minute) + timestamp[-3:]
            direction2.loc[timestamp, crossroadID] = direction2.loc[timestamp, crossroadID] + 1
pkl.dump([direction2],open(pwd+"/direction2.pkl",'wb'))



# 4, sum up directions of cross,   strore GPS in pred_cross['dir1','dir2',...,'dirNum']
for cross in pred_cross:
    for k in range(len(flow)):
        print(flow.iloc[k][cross])
        break





# 5, make data for model
#timeNoMap = {7:0,9:1,11:2,14:3,16:4,17:5} # No. of time_range
#time_start = datetime.strptime('2019-08-01 07:55:04','%Y-%m-%d %H:%M:%S')

timeNo_ = pd.DataFrame(np.zeros([133*12,1])) # 0-22, onehot-23, np.zeros([288*19,1])
flow_day_ = pd.DataFrame(np.zeros([133*12,7]))
flow_minute_ = pd.DataFrame(np.zeros([133*12,6]))
flow_y_ = pd.DataFrame(np.zeros([133*12,6]))
import copy
data_tr = {}

count = 0
tmp_day = list(range(7))
tmp_minute = list(range(6))
tmp_y = list(range(6))
for cross in pred_cross:
    count = 0
    timeNo = copy.copy(timeNo_)
    flow_day = copy.copy(flow_day_)
    flow_minute = copy.copy(flow_minute_)
    flow_y = copy.copy(flow_y_)
    for k in range(len(flow)):
        time_ = flow.index[k]
        time_tmp = int( int(time_.hour)*100 + int(time_.minute) )
        if time_.day >= 8 and time_.day <= 19 and time_tmp >= 730 and time_tmp <= 1830:
            #print(time_)
            timeNo.loc[count] = time_.hour*2 + (1 if time_.minute>=30 else 0) - 15
            for ii in range(1,8):
                #print(k-288*ii)
                tmp_day[ii-1] = flow.iloc[k-288*ii][cross]
            flow_day.loc[count] = tmp_day
            for ii in range(1,7):
                tmp_minute[ii-1] = flow.iloc[k-ii][cross]
            flow_minute.loc[count] = tmp_minute
            for ii in range(0,6):
                tmp_y[ii] = flow.iloc[k+ii][cross]
            flow_y.loc[count] = tmp_y
            count = count + 1
    data_tr[cross] = [timeNo,flow_day,flow_minute,flow_y]




# save key-variables
pkl.dump([flow,relate_cross,crossName,target_cross,pred_cross,similar_cross,roadNet,submitList,data_tr],open(pwd+"/data.pkl",'wb'))
